package com.fwmagic.spark.xgboost

import ml.dmlc.xgboost4j.scala.{Booster, XGBoost}
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.DoubleType


/**
 * Created by yihaibo on 2019-08-12.
 */
object TestXgb {
  def main(args: Array[String]): Unit = {
    implicit val  spark = SparkSession
      .builder()
      .appName("xgboost")
      .master("local[*]")
      .getOrCreate()

    val srcDF: DataFrame = spark.read.format("csv").option("header", "true")
      .option("timestampFormat", "yyyy/MM/dd HH:mm:ss ZZ")
      .load("data/xgboost/sample_xw.csv")

    srcDF.printSchema()
    srcDF.show(20,false)

    //转为double
    val df: DataFrame = srcDF.select(srcDF.columns.map {
      name => column(name).cast(DoubleType)
    }: _*)
    df.printSchema()
    df.show(false)

    val model: Booster = XGBoost.loadModel("data/xgboost/model_v3.model")
    val classification_model = new XGBoostClassificationModel(model.toString)

    val vectorAssembler = new VectorAssembler()
      .setInputCols(df.columns)
      .setOutputCol("features")
    val xgbInput: DataFrame = vectorAssembler.transform(df).select("features")

    xgbInput.show(false)

    val df2 = classification_model.transform(xgbInput)

    df2.show(false)

    df2.groupBy("prediction").count().show(false)

    spark.close()
  }
}
